#!/usr/bin/env python
# coding: utf-8
# copyRight by heibanke

import numpy as np
import time

# for hw_3-1
"""
两个向量a=[1,2,3],b=[1,1,1]
如何在a,b的平面上找到一个向量和a垂直。
"""


# for hw_3-2
def fourier(x, y, n):
    m = []
    for i in xrange(n):
        m.append(np.cos(x * i + x))
        m.append(np.sin(x * i + x))

    mx = np.mat(m).T

    y2 = y.reshape(y.shape[0], 1) - np.mean(y)

    w = np.linalg.inv(mx.T * mx) * mx.T * y2

    yw = np.array(mx * w) + np.mean(y)
    yw.shape = (yw.shape[0],)
    return yw


# for hw_3-3
from matplotlib import pyplot as plt
import matplotlib.animation as animation


def eigshow(A):
    w = np.linspace(0, 2 * np.pi, 100)
    x = np.array([np.cos(w), np.sin(w)])
    ax = A.dot(x)

    def update(num):
        s1.set_data(x[0, 0:num], x[1, 0:num])
        s2.set_data(ax[0, 0:num], ax[1, 0:num])
        return s1, s2

    fig1 = plt.figure()

    s1, = plt.plot([], [], 'ro')
    s2, = plt.plot([], [], 'bo')
    plt.xlim(-4, 4)
    plt.ylim(-4, 4)
    # plt.axis('equal')
    plt.xlabel('x')
    plt.title('Animation2')

    line_ani = animation.FuncAnimation(fig1, update, 100,
                                       interval=50, blit=False, repeat=False)
    # plt.axis('equal')
    plt.show()


"""
########################################################
#   下节课介绍内容
#
# for hw_3-4    
"""
import time


def time_cost(f):
    def _f(*arg, **kwarg):
        start = time.clock()
        a = f(*arg, **kwarg)
        end = time.clock()
        print f.__name__, "run cost time is ", end - start
        return a

    return _f


A = np.array([[1, 1], [1, 0]])
eigval, S = np.linalg.eig(A)
x = np.array([[1], [0]])
z = np.linalg.inv(S).dot(x)


@time_cost
def fib_eig_seq(seq):
    return [fib_eig(i) for i in seq]


def fib_eig(N):
    def lamda_n(lamda, n):
        a = np.array([[lamda[0] ** n, 0], [0, lamda[1] ** n]])
        return a

    if N >= 2:
        y = S.dot(lamda_n(eigval, N + 1).dot(z))
        return y[1]
    elif N < 0:
        assert False, "N must larger than zeros"
    else:
        return x[1 - N]


# for hw_3-5
"""
两个城市上海和北京，2020年人口为3200w和4000w，
虽然北京是中国的首都，不过由于雾霾的原因，每年留下0.8, 去上海0.2。
而上海留下0.9，去北京0.1。
问100年以后，人口怎样分布。1000年呢。（假设生老病死人数平衡，N年的总人数都不变）
"""


def hw_3_5(A, x, N_year):
    eigval, S = np.linalg.eig(A)

    Lamda = np.zeros(A.shape)
    for i in xrange(A.shape[0]):
        Lamda[i, i] = eigval[i] ** N_year

    z = S.dot(Lamda).dot(np.linalg.inv(S)).dot(x)

    print z


# for hw_3-6 
"""
手写字母的PCA抽取
"""


def pca(dataMat, k=0.9):
    covMat = np.cov(dataMat, rowvar=0)  # 求协方差方阵
    u, s, v = np.linalg.svd(covMat)

    for i in xrange(len(s) - 1):
        if sum(s[:i + 1]) / sum(s) > k:
            print i
            break
    u_s = u[:, :i + 1]
    lowd_sample = u_s.T.dot(dataMat.T)
    recon_sample = u_s.dot(lowd_sample)
    return recon_sample.T


def show_pic(sample, recon_sample):
    N = sample.shape[0]
    a = np.zeros((28 * N, 28 * 2))
    for i in xrange(N):
        a[28 * i:28 * (i + 1), :28] = sample[i, :].reshape(28, 28)
        a[28 * i:28 * (i + 1), 28:] = recon_sample[i, :].reshape(28, 28)
    plt.imshow(np.floor(a), cmap="gray")
    plt.show()
